#include "WaterShed.h"
#include <vector>

#define MARK_TEST \
	int markTest; \
	markTest = *((int *)(recordImageBits + bytesPerLine * pTest.y () + \
			pTest.x () * 4)); \
	if (markTest != 0) { \
		if (curMark != 0 && curMark != markTest) { \
			pDestImage->setPixel (p, qRgba (255, 255, 255, \
					qAlpha (pSrcImage->pixel (p)))); \
			graySortList[t].erase (graySortList[t].begin () + i); \
			--i; \
			continue; \
		} \
		pRecordImage->setPixel (p, markTest); \
		curMark = markTest; \
	} \
	if (qGray (pSrcImage->pixel (pTest)) <= qGray (pSrcImage->pixel (p))) \
		isLocalMin = false;


using namespace std;



WaterShed::WaterShed ()
{

}



Ptr< QImage > WaterShed::process (Ptr< QImage > pSrcImage)
{
	Ptr< QImage > pDestImage = new QImage (pSrcImage->size (),
			QImage::Format_ARGB32);
	pDestImage->fill (qRgba (0, 0, 0, 255));

	// Sort pixels
	vector< vector< QPoint > > graySortList (256);
	for (int v = 0; v < pSrcImage->height (); ++v) {
		QRgb * srcLine = (QRgb *)pSrcImage->scanLine (v);
		for (int u = 0; u < pSrcImage->width (); ++u) {
			graySortList[qGray(srcLine[u])].push_back (QPoint (u, v));
		}
	}

	Ptr< QImage > pRecordImage = new QImage (pSrcImage->size (),
			QImage::Format_ARGB32);
	pRecordImage->fill (0);
	int mark = 1;
	int bytesPerLine = pRecordImage->bytesPerLine ();
	unsigned char * recordImageBits = pRecordImage->bits ();
	unsigned char * destImageBits = pDestImage->bits ();
	int size;
	for (int t = 0; t < 256; ++t) {
		size = 0;
		while (graySortList[t].empty () == false) {
			if (size == graySortList[t].size ()) {
				pRecordImage->setPixel (graySortList[t][0], mark);
				++mark;
				graySortList[t].erase (graySortList[t].begin ());
			}
			size = graySortList[t].size ();
			for (int i = 0; i < graySortList[t].size (); ++i) {
				QPoint p = graySortList[t][i];
				int curMark = 0;
				bool isLocalMin = true;
				// Left
				if (p.x () - 1 >= 0) {
					QPoint pTest = QPoint (p.x () - 1, p.y ());
					int markTest;
						markTest = *((int *)(recordImageBits + bytesPerLine * pTest.y () +
						pTest.x () * 4));
						if (markTest != 0) {
							if (curMark != 0 && curMark != markTest) {
								pDestImage->setPixel (p, qRgba (255, 255, 255,
								qAlpha (pSrcImage->pixel (p))));
								graySortList[t].erase (graySortList[t].begin () + i);
								--i;
								continue;
							}
							pRecordImage->setPixel (p, markTest);
							curMark = markTest;
						}
						if (qGray (pSrcImage->pixel (pTest)) <= qGray (pSrcImage->pixel (p))) \
							isLocalMin = false;
				}
				// Right
				if (p.x () + 1 < pRecordImage->width ()) {
					QPoint pTest = QPoint (p.x () + 1, p.y ());
					int markTest;
					markTest = *((int *)(recordImageBits + bytesPerLine * pTest.y () +
						pTest.x () * 4));
					if (markTest != 0) {
						if (curMark != 0 && curMark != markTest) {
							pDestImage->setPixel (p, qRgba (255, 255, 255,
								qAlpha (pSrcImage->pixel (p))));
							graySortList[t].erase (graySortList[t].begin () + i);
							--i;
							continue;
						}
						pRecordImage->setPixel (p, markTest);
						curMark = markTest;
					}
					if (qGray (pSrcImage->pixel (pTest)) <= qGray (pSrcImage->pixel (p))) \
						isLocalMin = false;
				}
				// Top
				if (p.y () - 1 >= 0) {
					QPoint pTest = QPoint (p.x (), p.y () - 1);
					int markTest;
					markTest = *((int *)(recordImageBits + bytesPerLine * pTest.y () +
						pTest.x () * 4));
					if (markTest != 0) {
						if (curMark != 0 && curMark != markTest) {
							pDestImage->setPixel (p, qRgba (255, 255, 255,
								qAlpha (pSrcImage->pixel (p))));
							graySortList[t].erase (graySortList[t].begin () + i);
							--i;
							continue;
						}
						pRecordImage->setPixel (p, markTest);
						curMark = markTest;
					}
					if (qGray (pSrcImage->pixel (pTest)) <= qGray (pSrcImage->pixel (p))) \
						isLocalMin = false;
				}
				// Bottom
				if (p.y () + 1 < pRecordImage->height ()) {
					QPoint pTest = QPoint (p.x (), p.y () + 1);
					int markTest;
					markTest = *((int *)(recordImageBits + bytesPerLine * pTest.y () +
						pTest.x () * 4));
					if (markTest != 0) {
						if (curMark != 0 && curMark != markTest) {
							pDestImage->setPixel (p, qRgba (255, 255, 255,
								qAlpha (pSrcImage->pixel (p))));
							graySortList[t].erase (graySortList[t].begin () + i);
							--i;
							continue;
						}
						pRecordImage->setPixel (p, markTest);
						curMark = markTest;
					}
					if (qGray (pSrcImage->pixel (pTest)) <= qGray (pSrcImage->pixel (p))) \
						isLocalMin = false;
				}
				if (isLocalMin) {
					pRecordImage->setPixel (p, mark);
					++mark;
					graySortList[t].erase (graySortList[t].begin () + i);
					--i;
				} else if (curMark != 0) {
					graySortList[t].erase (graySortList[t].begin () + i);
					--i;
				}
			}
		}
	}
	for (int v = 0; v < pRecordImage->height (); ++v) {
		QRgb * recordLine = (QRgb *)pRecordImage->scanLine (v);
		for (int u = 0; u < pRecordImage->width (); ++u) {
			recordLine[u] |= 0xFF000000;
		}
	}
	return pDestImage;
}
